import torch
import torch.nn as nn
import matplotlib.pyplot as plt


def test():
    # 读取图像, 形状: (640, 640, 3)
    img = plt.imread('data/img.jpg')
    plt.imshow(img)
    plt.axis('off')
    plt.show()
    # 构建卷积层
    # out_channels表示卷积核个数
    # 修改out_channels，stride，padding观察特征图的变化情况
    conv = nn.Conv2d(in_channels=3, out_channels=3, kernel_size=3, stride=2, padding=0)
    # 输入形状: (BatchSize, Channel, Height, Width)
    # img形状: torch.Size([3, 640, 640])
    img = torch.tensor(img).permute(2, 0, 1)
    # img形状: torch.Size([1, 3, 640, 640])
    img = img.unsqueeze(0)
    # 将图像送入卷积层中
    feature_map_img = conv(img.to(torch.float32))
    # 打印特征图的形状
    print(feature_map_img.shape)


if __name__ == '__main__':
    test()
